import os
import yaml
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.utils.data import Dataset, DataLoader
import musdb


class MUSDBDataset(Dataset):
    def __init__(self, root, subset='train', segment_length=44100):
        self.mus = musdb.DB(root=root, subsets=[subset], is_wav=True)
        self.segment_length = segment_length

    def __len__(self):
        return len(self.mus)

    def __getitem__(self, idx):
        track = self.mus.tracks[idx]
        mixture = track.audio.T.astype(np.float32)  # (2, T)
        targets = np.stack([track.targets[name].audio.T for name in sorted(track.sources)], axis=0)

        # 随机截取片段
        if mixture.shape[1] > self.segment_length:
            start = np.random.randint(0, mixture.shape[1] - self.segment_length)
            mixture = mixture[:, start:start + self.segment_length]
            targets = targets[:, :, start:start + self.segment_length]

        return torch.from_numpy(mixture), torch.from_numpy(targets)


def save_checkpoint(model, optimizer, epoch, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)


def load_config(config_path):
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def plot_training_curves(log_path, save_path):
    log = np.load(log_path)
    plt.figure(figsize=(12, 6))
    plt.plot(log['train_loss'], label='Train Loss')
    plt.plot(log['valid_loss'], label='Validation Loss')
    plt.title('Training Curves')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(save_path)
    plt.close()